
import cv2
import numpy as np
import matplotlib.pyplot as plt


def main(args):
    if args.gray:
        img = cv2.imread(args.src, cv2.IMREAD_GRAYSCALE)
    else:
        img = cv2.imread(args.src, cv2.IMREAD_COLOR)
    print(img.dtype, img.shape)

    # cv.imshow('input', img)

    # res = cv2.equalizeHist(img)
    # cv2.imshow('equalizeHist', res)
    # cv2.waitKey(0)

    # plt.hist(
    #     img.ravel(),
    #     bins=256, range=[0, 256],
    #     density=True,
    # )

    if args.gray:
        hist = cv2.calcHist(images=[img], channels=[0], mask=None, histSize=[256], ranges=(0, 256))
    else:
        hist = cv2.calcHist(images=[img], channels=[0, 1, 2], mask=None, histSize=[256, 256, 256], ranges=(0, 256, 0, 256, 0, 256))
    # print(len(hist), hist)

    area = np.cumsum(hist)
    # print(len(area), area)

    xs = np.percentile(img, [
        0.05,
        0.10,
        0.15,
        0.20,
        99.80,
        99.85,
        99.90,
        99.95,
    ])
    print('lining', xs)

    plt.vlines(xs, 0, area.max(), linestyles='dashed', colors='red')
    plt.hlines(hist.max(), 0, 256, linestyles='dashed', colors='red')
    plt.hlines(area.max(), 0, 256, linestyles='dashed', colors='red')
    plt.plot(hist, color='black')
    plt.plot(area, color='blue')
    plt.xlim([0, 256])
    plt.show()

    # 一维像素直方图，也即是单通道直方图
    # color = ('blue', 'green', 'red')
    # for i, color in enumerate(color):
    #     hist = cv2.calcHist([img], channels=[i], mask=None, histSize=[256], range=[0, 256])
    #     print(hist)
    #     plt.plot(hist, color=color)
    #     plt.xlim([0, 256])
    # plt.show()

    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src')
    parser.add_argument('--gray', action='store_true')
    args = parser.parse_args()
    main(args)
